{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "np.set_printoptions(threshold=np.inf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k_input = \"input/input_k.npy\"\n",
    "v_input = \"input/input_v.npy\"\n",
    "w_input = \"input/input_w.npy\"\n",
    "r_input = \"input/input_q.npy\"\n",
    "a_input = \"input/input_a.npy\"\n",
    "b_input = \"input/input_b.npy\"\n",
    "h_input = \"input/input_h0.npy\"\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = np.load(k_input) \n",
    "v = np.load(v_input) \n",
    "w = np.load(w_input) \n",
    "q = np.load(r_input) \n",
    "a = np.load(a_input) \n",
    "b = np.load(b_input) \n",
    "h = np.load(h_input) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "B = 1\n",
    "T = 16\n",
    "N = 64\n",
    "H = 1\n",
    "data_type = np.float32\n",
    "param_shape = (B, H, T, N)\n",
    "u_shape = (H, N)\n",
    "state_shape = (B, H, N, N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h = np.zeros(state_shape).astype(data_type)\n",
    "o = np.zeros(param_shape).astype(data_type)\n",
    "for t in range(T):\n",
    "        for bi in range(B):\n",
    "            for hi in range(H):\n",
    "                q_t = q[bi, hi, t]\n",
    "                k_t = k[bi, hi, t] * 0.5\n",
    "                v_t = v[bi, hi, t] * 0.5\n",
    "                a_t = a[bi, hi, t]\n",
    "                b_t = b[bi, hi, t]\n",
    "                w_t = np.exp(-np.exp(w[bi, hi, t]))\n",
    "\n",
    "                aa = a_t[None, :]\n",
    "                saa = a_t[None, :] * h[bi, hi]\n",
    "                sa = np.sum(saa, axis=1)\n",
    "                bb = b_t[None, :]\n",
    "                sab = sa * b_t[None, :]\n",
    "                tmp1 = h[bi, hi] * w_t[None, :]\n",
    "                tmp2 = tmp1 + k_t[None, :] * v_t[:, None]\n",
    "                tmp3 = tmp2 + sab\n",
    "                h[bi, hi] = tmp3   \n",
    "                y = np.sum((h[bi, hi] * q_t[None, :]), axis=1)\n",
    "\n",
    "                o[bi, hi, t] = y\n",
    "            print(\"aa\", aa)\n",
    "            print(\"bb\", bb)\n",
    "            print(\"sa\", sa)\n",
    "            print(\"sab\", sab)\n",
    "            print(\"tmp1\", tmp1)\n",
    "            print(\"tmp2\", tmp2)\n",
    "            print(\"tmp3\", tmp3)\n",
    "            print(\"h\", h)\n",
    "            print(\"o\", o)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# h = np.zeros(state_shape).astype(data_type)\n",
    "# o = np.zeros(param_shape).astype(data_type)\n",
    "# for t in range(T):\n",
    "#         for bi in range(B):\n",
    "#             for hi in range(H):\n",
    "#                 q_t = q[bi, hi, t]\n",
    "#                 k_t = k[bi, hi, t] * 0.5\n",
    "#                 v_t = v[bi, hi, t] * 0.5\n",
    "#                 a_t = a[bi, hi, t]\n",
    "#                 b_t = b[bi, hi, t]\n",
    "#                 w_t = np.exp(-np.exp(w[bi, hi, t]))\n",
    "\n",
    "#                 sa = np.sum((a_t[None, :] * h[bi, hi]), axis=1)\n",
    "\n",
    "#                 h[bi, hi] = (h[bi, hi] * w_t[None, :] +  # [N,V] * [N,1]\n",
    "#                                  k_t[None, :] * v_t[:, None] +     # [N,1] * [1,V]\n",
    "#                                  sa * b_t[None, :])                # [V] * [N,1]\n",
    "\n",
    "#                 y = np.sum((h[bi, hi] * q_t[None, :]), axis=1)\n",
    "\n",
    "#                 o[bi, hi, t] = y\n",
    "#             # print(\"sa\", sa)\n",
    "#             print(\"h\", h)\n",
    "#             print(\"o\", o)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# h = np.zeros(state_shape).astype(data_type)\n",
    "# o = np.zeros(param_shape).astype(data_type)\n",
    "# for t in range(T):\n",
    "#         for bi in range(B):\n",
    "#             for hi in range(H):\n",
    "#                 q_t = q[bi, hi, t]\n",
    "#                 k_t = k[bi, hi, t] * 0.5\n",
    "#                 v_t = v[bi, hi, t] * 0.5\n",
    "#                 a_t = a[bi, hi, t]\n",
    "#                 b_t = b[bi, hi, t]\n",
    "#                 w_t = np.exp(-np.exp(w[bi, hi, t]))\n",
    "\n",
    "#                 sa = np.sum((a_t[:, None] * h[bi, hi]), axis=0)\n",
    "                \n",
    "#                 tmp1 = h[bi, hi] * w_t[:, None]\n",
    "#                 kv = k_t[:, None] * v_t[None, :]\n",
    "#                 sab = sa * b_t[:, None]\n",
    "#                 h[bi, hi] = (tmp1 +  # [N,V] * [N,1]\n",
    "#                                 kv  +     # [N,1] * [1,V]\n",
    "#                             sab )             # [V] * [N,1]\n",
    "\n",
    "#                 y = np.sum((h[bi, hi] * q_t[:, None]), axis=0)\n",
    "\n",
    "#                 o[bi, hi, t] = y\n",
    "#             # print(\"tmp1\", tmp1)\n",
    "#             # print(\"kv\", kv)\n",
    "#             # print(\"sab\", sab)\n",
    "#             print(\"h\", h)\n",
    "#             print(\"o\",o)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MindSpore",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
